{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a13ddc8",
   "metadata": {},
   "source": [
    "# PromptTools : LLM Output Evaluation using LanceDB\n",
    "![picture](https://media.licdn.com/dms/image/D5622AQG7LlHPaun5oQ/feedshare-shrink_2048_1536/0/1692934644617?e=1700092800&v=beta&t=PIHTwJKBDFGzgyTINYWq-saXgIkC37AvEkgMFf3SbK0)\n",
    "\n",
    "Find more on the <a href=\"https://hegel-ai.com/\">Prompttools</a> page."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "623f0cfe",
   "metadata": {},
   "source": [
    "## Installations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "723d11f0-2426-42a1-a122-960840859e68",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet --force-reinstall prompttools lancedb sentence_transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "622dea9a",
   "metadata": {},
   "source": [
    "## Run an experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5650e31",
   "metadata": {},
   "source": [
    "One common use case is to compare **two different embedding functions** and how it may impact your document retrieval. We have can define what embedding functions we'd like to test here.\n",
    "\n",
    "Note: If you previously haven't downloaded these embedding models. This may kick off downloads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "821bbb21-292c-44e5-bdf0-ab05350acb36",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import os\n",
    "from prompttools.experiment import LanceDBExperiment\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "# Configuring the environment variable OPENAI_API_KEY\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    # OR set the key here as a variable\n",
    "    openai.api_key = \"...\"\n",
    "\n",
    "\n",
    "DEFAULT = SentenceTransformer(\"paraphrase-MiniLM-L3-v2\")\n",
    "MIMNILM_L6 = SentenceTransformer(\"all-MiniLM-L6-v2\")\n",
    "\n",
    "\n",
    "def default_embed_func(batch):\n",
    "    return [DEFAULT.encode(sentence) for sentence in batch]\n",
    "\n",
    "\n",
    "def minilm_l6_embed_func(batch):\n",
    "    return [MIMNILM_L6.encode(sentence) for sentence in batch]\n",
    "\n",
    "\n",
    "def openai_ada2_embed_func(batch):\n",
    "    rs = openai.Embedding.create(input=batch, engine=\"text-embedding-ada-002\")\n",
    "    return [record[\"embedding\"] for record in rs[\"data\"]]\n",
    "\n",
    "\n",
    "emb_fns = {\"minilm_l6\": minilm_l6_embed_func, \"default\": default_embed_func}\n",
    "# Try with openai\n",
    "# emb_fns = {\"openai-ada-002\": openai_ada2_embed_func, \"minilm_l6\": minilm_l6_embed_func,  \"default\": default_embed_func }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "189752e1-f392-42a3-b21c-56e8b384d9fc",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9114cfbf",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-07-20T07:13:56.829960Z",
     "start_time": "2023-07-20T07:13:56.825481Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "use_existing_table = (\n",
    "    False  # Specify that we want to create a collection during the experiment\n",
    ")\n",
    "\n",
    "# Documents that will be added into the database. LanceDB also accepts other dataset formats like pydict, pyarrow, Pydantic etc.\n",
    "# Learn more here - https://lancedb.github.io/lancedb/guides/tables/\n",
    "\n",
    "data = pd.DataFrame(\n",
    "    {\n",
    "        \"text\": [\n",
    "            \"This is a document\",\n",
    "            \"This is another document\",\n",
    "            \"This is the document.\",\n",
    "        ],\n",
    "        \"metadatas\": [\n",
    "            {\"source\": \"my_source\"},\n",
    "            {\"source\": \"my_source\"},\n",
    "            {\"source\": \"my_source\"},\n",
    "        ],\n",
    "        \"ids\": [\"id1\", \"id2\", \"id3\"],\n",
    "    }\n",
    ")\n",
    "\n",
    "query_args = {\n",
    "    \"text\": [\"This is a query document\", \"This is a another query document\"],\n",
    "    \"metric\": [\"cosine\", \"l2\"],\n",
    "}\n",
    "\n",
    "\n",
    "# Set up the experiment\n",
    "experiment = LanceDBExperiment(\n",
    "    data=data,\n",
    "    embedding_fns=emb_fns,\n",
    "    query_args=query_args,\n",
    ")\n",
    "\n",
    "# [Optional] Advanced query args\n",
    "# Our test queries, along with optional query args. LanceDB query accepts a few args to customize your search:\n",
    "# metrics: \"l2\", \"cosine\", or \"dot\" (cosine by default)\n",
    "# filter: SQL where clause to filter the vector search results before applying the limit. (None by default)\n",
    "# limit: number of results to return (3 by default)\n",
    "\"\"\"\n",
    "query_args_adv = {\n",
    "                \"text\": [\"This is a query document\", \"This is a another query document\"], \n",
    "                \"metric\": [\"cosine\", \"l2\", \"dot\"],\n",
    "                \"filter\": [\"text IS NOT NULL\" , \"text LIKE '%document.%'\"]\n",
    "                }\n",
    "experiment = LanceDBExperiment(\n",
    "    data=data,\n",
    "    embedding_fns=emb_fns,\n",
    "    query_args=query_args_adv,\n",
    " \n",
    ")\n",
    "\"\"\"\n",
    "print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3fa5450",
   "metadata": {},
   "source": [
    "We can then run the experiment to get results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "83b33130",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-07-20T07:16:21.469371Z",
     "start_time": "2023-07-20T07:16:21.462342Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: rate limit only support up to 3.10, proceeding without rate limiter\n",
      "WARNING: rate limit only support up to 3.10, proceeding without rate limiter\n",
      "WARNING: rate limit only support up to 3.10, proceeding without rate limiter\n"
     ]
    }
   ],
   "source": [
    "experiment.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b013dca",
   "metadata": {},
   "source": [
    "We can visualize the result. In this case, the result of the second query \"This is a another query document\" is different.\n",
    "\n",
    "paraphrase-MiniLM-L3-v2: [id2, id3, id1]\n",
    "\n",
    "default (all-MiniLM-L6-v2) : [id2, id1, id3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4908a9e3-3421-43f8-8a55-df9b272c7a7e",
   "metadata": {},
   "source": [
    "## Let's visualize the outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "01c7e682",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>metric</th>\n",
       "      <th>emb_fn</th>\n",
       "      <th>top doc ids</th>\n",
       "      <th>distances</th>\n",
       "      <th>documents</th>\n",
       "      <th>latency</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>openai-ada-002</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[0.06923848390579224, 0.08619403839111328, 0.10083281993865967]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>2.047757</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>openai-ada-002</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[0.1383669227361679, 0.17228886485099792, 0.2016243040561676]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.392724</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>openai-ada-002</td>\n",
       "      <td>[id2, id1, id3]</td>\n",
       "      <td>[0.05721437931060791, 0.08627921342849731, 0.1029735803604126]</td>\n",
       "      <td>[This is another document, This is a document, This is the document.]</td>\n",
       "      <td>5.544514</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>openai-ada-002</td>\n",
       "      <td>[id2, id1, id3]</td>\n",
       "      <td>[0.11442875862121582, 0.17255812883377075, 0.20594733953475952]</td>\n",
       "      <td>[This is another document, This is a document, This is the document.]</td>\n",
       "      <td>0.622729</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>minilm_l6</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[0.8099705576896667, 0.8289484977722168, 0.8308900594711304]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.198556</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>minilm_l6</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[1.619940996170044, 1.6578971147537231, 1.6617801189422607]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.190643</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>minilm_l6</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[0.8099705576896667, 0.8289484977722168, 0.8308900594711304]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.267542</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>minilm_l6</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[1.619940996170044, 1.6578971147537231, 1.6617801189422607]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.262355</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>default</td>\n",
       "      <td>[id2, id3, id1]</td>\n",
       "      <td>[0.7633732557296753, 0.773878812789917, 0.7882261872291565]</td>\n",
       "      <td>[This is another document, This is the document., This is a document]</td>\n",
       "      <td>0.111481</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>default</td>\n",
       "      <td>[id3, id1, id2]</td>\n",
       "      <td>[45.84406280517578, 49.12738037109375, 49.839256286621094]</td>\n",
       "      <td>[This is the document., This is a document, This is another document]</td>\n",
       "      <td>0.113457</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>default</td>\n",
       "      <td>[id2, id3, id1]</td>\n",
       "      <td>[0.7633732557296753, 0.773878812789917, 0.7882261872291565]</td>\n",
       "      <td>[This is another document, This is the document., This is a document]</td>\n",
       "      <td>0.138935</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>default</td>\n",
       "      <td>[id3, id1, id2]</td>\n",
       "      <td>[45.84406280517578, 49.12738037109375, 49.839256286621094]</td>\n",
       "      <td>[This is the document., This is a document, This is another document]</td>\n",
       "      <td>0.139161</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                text  metric          emb_fn      top doc ids  \\\n",
       "0   This is a query document          cosine  openai-ada-002  [id1, id3, id2]   \n",
       "1   This is a query document          l2      openai-ada-002  [id1, id3, id2]   \n",
       "2   This is a another query document  cosine  openai-ada-002  [id2, id1, id3]   \n",
       "3   This is a another query document  l2      openai-ada-002  [id2, id1, id3]   \n",
       "4   This is a query document          cosine  minilm_l6       [id1, id3, id2]   \n",
       "5   This is a query document          l2      minilm_l6       [id1, id3, id2]   \n",
       "6   This is a another query document  cosine  minilm_l6       [id1, id3, id2]   \n",
       "7   This is a another query document  l2      minilm_l6       [id1, id3, id2]   \n",
       "8   This is a query document          cosine  default         [id2, id3, id1]   \n",
       "9   This is a query document          l2      default         [id3, id1, id2]   \n",
       "10  This is a another query document  cosine  default         [id2, id3, id1]   \n",
       "11  This is a another query document  l2      default         [id3, id1, id2]   \n",
       "\n",
       "                                                          distances  \\\n",
       "0   [0.06923848390579224, 0.08619403839111328, 0.10083281993865967]   \n",
       "1   [0.1383669227361679, 0.17228886485099792, 0.2016243040561676]     \n",
       "2   [0.05721437931060791, 0.08627921342849731, 0.1029735803604126]    \n",
       "3   [0.11442875862121582, 0.17255812883377075, 0.20594733953475952]   \n",
       "4   [0.8099705576896667, 0.8289484977722168, 0.8308900594711304]      \n",
       "5   [1.619940996170044, 1.6578971147537231, 1.6617801189422607]       \n",
       "6   [0.8099705576896667, 0.8289484977722168, 0.8308900594711304]      \n",
       "7   [1.619940996170044, 1.6578971147537231, 1.6617801189422607]       \n",
       "8   [0.7633732557296753, 0.773878812789917, 0.7882261872291565]       \n",
       "9   [45.84406280517578, 49.12738037109375, 49.839256286621094]        \n",
       "10  [0.7633732557296753, 0.773878812789917, 0.7882261872291565]       \n",
       "11  [45.84406280517578, 49.12738037109375, 49.839256286621094]        \n",
       "\n",
       "                                                                documents  \\\n",
       "0   [This is a document, This is the document., This is another document]   \n",
       "1   [This is a document, This is the document., This is another document]   \n",
       "2   [This is another document, This is a document, This is the document.]   \n",
       "3   [This is another document, This is a document, This is the document.]   \n",
       "4   [This is a document, This is the document., This is another document]   \n",
       "5   [This is a document, This is the document., This is another document]   \n",
       "6   [This is a document, This is the document., This is another document]   \n",
       "7   [This is a document, This is the document., This is another document]   \n",
       "8   [This is another document, This is the document., This is a document]   \n",
       "9   [This is the document., This is a document, This is another document]   \n",
       "10  [This is another document, This is the document., This is a document]   \n",
       "11  [This is the document., This is a document, This is another document]   \n",
       "\n",
       "     latency  \n",
       "0   2.047757  \n",
       "1   0.392724  \n",
       "2   5.544514  \n",
       "3   0.622729  \n",
       "4   0.198556  \n",
       "5   0.190643  \n",
       "6   0.267542  \n",
       "7   0.262355  \n",
       "8   0.111481  \n",
       "9   0.113457  \n",
       "10  0.138935  \n",
       "11  0.139161  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "experiment.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "266c13eb",
   "metadata": {},
   "source": [
    "## Evaluate the model response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bebb8023",
   "metadata": {},
   "source": [
    "To evaluate the results, we'll define an **evaluation function**. \n",
    "Sometimes, you know order of the most relevant document should be given a query, and you can compute the correlation between expected ranking and actual ranking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8ddbb951",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.stats as stats\n",
    "\n",
    "# For each query, you can define what the expected ranking is.\n",
    "EXPECTED_RANKING = {\n",
    "    \"This is a query document\": [\"id1\", \"id3\", \"id2\"],\n",
    "    \"This is a another query document\": [\"id2\", \"id3\", \"id1\"],\n",
    "}\n",
    "\n",
    "\n",
    "def measure_correlation(\n",
    "    row: \"pandas.core.series.Series\", ranking_column_name: str = \"top doc ids\"\n",
    ") -> float:\n",
    "    r\"\"\"\n",
    "    A simple test that compares the expected ranking for a given query with the actual ranking produced\n",
    "    by the embedding function being tested.\n",
    "    \"\"\"\n",
    "    input_query = row[\"text\"]\n",
    "    correlation, _ = stats.spearmanr(\n",
    "        row[ranking_column_name], EXPECTED_RANKING[input_query]\n",
    "    )\n",
    "    return correlation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "974d6065",
   "metadata": {},
   "source": [
    "## Finally, we can evaluate and visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e80dfeec",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "experiment.evaluate(\"ranking_correlation\", measure_correlation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4d09c18e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>metric</th>\n",
       "      <th>emb_fn</th>\n",
       "      <th>top doc ids</th>\n",
       "      <th>distances</th>\n",
       "      <th>documents</th>\n",
       "      <th>latency</th>\n",
       "      <th>ranking_correlation</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>openai-ada-002</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[0.06923848390579224, 0.08619403839111328, 0.10083281993865967]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>2.047757</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>openai-ada-002</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[0.1383669227361679, 0.17228886485099792, 0.2016243040561676]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.392724</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>openai-ada-002</td>\n",
       "      <td>[id2, id1, id3]</td>\n",
       "      <td>[0.05721437931060791, 0.08627921342849731, 0.1029735803604126]</td>\n",
       "      <td>[This is another document, This is a document, This is the document.]</td>\n",
       "      <td>5.544514</td>\n",
       "      <td>-1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>openai-ada-002</td>\n",
       "      <td>[id2, id1, id3]</td>\n",
       "      <td>[0.11442875862121582, 0.17255812883377075, 0.20594733953475952]</td>\n",
       "      <td>[This is another document, This is a document, This is the document.]</td>\n",
       "      <td>0.622729</td>\n",
       "      <td>-1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>minilm_l6</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[0.8099705576896667, 0.8289484977722168, 0.8308900594711304]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.198556</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>minilm_l6</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[1.619940996170044, 1.6578971147537231, 1.6617801189422607]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.190643</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>minilm_l6</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[0.8099705576896667, 0.8289484977722168, 0.8308900594711304]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.267542</td>\n",
       "      <td>0.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>minilm_l6</td>\n",
       "      <td>[id1, id3, id2]</td>\n",
       "      <td>[1.619940996170044, 1.6578971147537231, 1.6617801189422607]</td>\n",
       "      <td>[This is a document, This is the document., This is another document]</td>\n",
       "      <td>0.262355</td>\n",
       "      <td>0.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>default</td>\n",
       "      <td>[id2, id3, id1]</td>\n",
       "      <td>[0.7633732557296753, 0.773878812789917, 0.7882261872291565]</td>\n",
       "      <td>[This is another document, This is the document., This is a document]</td>\n",
       "      <td>0.111481</td>\n",
       "      <td>0.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>This is a query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>default</td>\n",
       "      <td>[id3, id1, id2]</td>\n",
       "      <td>[45.84406280517578, 49.12738037109375, 49.839256286621094]</td>\n",
       "      <td>[This is the document., This is a document, This is another document]</td>\n",
       "      <td>0.113457</td>\n",
       "      <td>-1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>cosine</td>\n",
       "      <td>default</td>\n",
       "      <td>[id2, id3, id1]</td>\n",
       "      <td>[0.7633732557296753, 0.773878812789917, 0.7882261872291565]</td>\n",
       "      <td>[This is another document, This is the document., This is a document]</td>\n",
       "      <td>0.138935</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>This is a another query document</td>\n",
       "      <td>l2</td>\n",
       "      <td>default</td>\n",
       "      <td>[id3, id1, id2]</td>\n",
       "      <td>[45.84406280517578, 49.12738037109375, 49.839256286621094]</td>\n",
       "      <td>[This is the document., This is a document, This is another document]</td>\n",
       "      <td>0.139161</td>\n",
       "      <td>-0.5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                text  metric          emb_fn      top doc ids  \\\n",
       "0   This is a query document          cosine  openai-ada-002  [id1, id3, id2]   \n",
       "1   This is a query document          l2      openai-ada-002  [id1, id3, id2]   \n",
       "2   This is a another query document  cosine  openai-ada-002  [id2, id1, id3]   \n",
       "3   This is a another query document  l2      openai-ada-002  [id2, id1, id3]   \n",
       "4   This is a query document          cosine  minilm_l6       [id1, id3, id2]   \n",
       "5   This is a query document          l2      minilm_l6       [id1, id3, id2]   \n",
       "6   This is a another query document  cosine  minilm_l6       [id1, id3, id2]   \n",
       "7   This is a another query document  l2      minilm_l6       [id1, id3, id2]   \n",
       "8   This is a query document          cosine  default         [id2, id3, id1]   \n",
       "9   This is a query document          l2      default         [id3, id1, id2]   \n",
       "10  This is a another query document  cosine  default         [id2, id3, id1]   \n",
       "11  This is a another query document  l2      default         [id3, id1, id2]   \n",
       "\n",
       "                                                          distances  \\\n",
       "0   [0.06923848390579224, 0.08619403839111328, 0.10083281993865967]   \n",
       "1   [0.1383669227361679, 0.17228886485099792, 0.2016243040561676]     \n",
       "2   [0.05721437931060791, 0.08627921342849731, 0.1029735803604126]    \n",
       "3   [0.11442875862121582, 0.17255812883377075, 0.20594733953475952]   \n",
       "4   [0.8099705576896667, 0.8289484977722168, 0.8308900594711304]      \n",
       "5   [1.619940996170044, 1.6578971147537231, 1.6617801189422607]       \n",
       "6   [0.8099705576896667, 0.8289484977722168, 0.8308900594711304]      \n",
       "7   [1.619940996170044, 1.6578971147537231, 1.6617801189422607]       \n",
       "8   [0.7633732557296753, 0.773878812789917, 0.7882261872291565]       \n",
       "9   [45.84406280517578, 49.12738037109375, 49.839256286621094]        \n",
       "10  [0.7633732557296753, 0.773878812789917, 0.7882261872291565]       \n",
       "11  [45.84406280517578, 49.12738037109375, 49.839256286621094]        \n",
       "\n",
       "                                                                documents  \\\n",
       "0   [This is a document, This is the document., This is another document]   \n",
       "1   [This is a document, This is the document., This is another document]   \n",
       "2   [This is another document, This is a document, This is the document.]   \n",
       "3   [This is another document, This is a document, This is the document.]   \n",
       "4   [This is a document, This is the document., This is another document]   \n",
       "5   [This is a document, This is the document., This is another document]   \n",
       "6   [This is a document, This is the document., This is another document]   \n",
       "7   [This is a document, This is the document., This is another document]   \n",
       "8   [This is another document, This is the document., This is a document]   \n",
       "9   [This is the document., This is a document, This is another document]   \n",
       "10  [This is another document, This is the document., This is a document]   \n",
       "11  [This is the document., This is a document, This is another document]   \n",
       "\n",
       "     latency  ranking_correlation  \n",
       "0   2.047757  1.0                  \n",
       "1   0.392724  1.0                  \n",
       "2   5.544514 -1.0                  \n",
       "3   0.622729 -1.0                  \n",
       "4   0.198556  1.0                  \n",
       "5   0.190643  1.0                  \n",
       "6   0.267542  0.5                  \n",
       "7   0.262355  0.5                  \n",
       "8   0.111481  0.5                  \n",
       "9   0.113457 -1.0                  \n",
       "10  0.138935  1.0                  \n",
       "11  0.139161 -0.5                  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "experiment.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaf4a45b",
   "metadata": {},
   "source": [
    "You can also use auto evaluation. We will add an example of this in the near future."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
